{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FINN - Transformation passes\n",
    "--------------------------------------\n",
    "In this notebook the idea behind transformation passes in FINN will be explained and with the help of an example the procedure of a transformation will be shown.\n",
    "\n",
    "We'll use the following utility functions to print the source code for function calls (`showSrc()`) and to visualize a network using netron (`showInNetron()`) in the Jupyter notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from finn.util.visualization import showSrc, showInNetron"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## General Information\n",
    "-----------------------------\n",
    "* changes (transforms) the given graph\n",
    "* input: ModelWrapper\n",
    "* returns the changed model (ModelWrapper) and flag `model_was_changed`\n",
    "\n",
    "Transformation passes have a base class and must inherit from that. Transformations are meant to be applied using .transform function from the ModelWrapper. This function makes a deep copy of the input model by default. The next cell shows .transform of ModelWrapper.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### .transform() from ModelWrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qonnx.core.modelwrapper import ModelWrapper\n",
    "showSrc(ModelWrapper.transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When the function is called, the model, the name of the transformation and, if required, the flag make_deepcopy are passed. It is also possible not to make a copy of the model. In this case `make_deepcopy` must be set to False. Then the branch `if make_deepcopy:` would not be taken and no copy of the model would be made. \n",
    "\n",
    "Additionally, the attribute `fix_float64` of the model is checked and if it is set to `True` all double values are converted to float. This assures a correct functionality of the model.\n",
    "\n",
    "The unchanged model is passed to the variable `transformed_model` to pass this variable on to the transformation later. \n",
    "\n",
    "`model_was_changed` indicates whether the transformation needs to be applied more than once. Because it needs to be applied at least one time `model_was_changed` is first set to True and then depending on the return values of the transformation function the transformation can be applied more then once. \n",
    "\n",
    "**Important**: Due to the structure of this function, `model_was_changed` must be set to False at some point. Otherwise the loop is infinite.\n",
    "    \n",
    "\n",
    "Each new transformation must correspond to the scheme of the base class and contain at least the function `apply(model)`, which returns the changed model and a bool value for `model_was_changed`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transformation Base Class     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qonnx.transformation.base import Transformation\n",
    "\n",
    "showSrc(Transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Base class is abstract class (`import ABC`) with only one abstract method (`apply()`) which gets the model as input. To show what a transformation should look like, the following example is taken from FINN."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example - ConvertSubToAdd\n",
    "-----------------------------\n",
    "The transformation replaces all subtraction nodes in a model with addition nodes with appropriate sign. For that an onnx model is loaded which contains one subtraction node. \n",
    "    \n",
    "Netron is used to visualize the result before and after."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "notebook_dir = os.environ['FINN_ROOT'] + \"/notebooks\"\n",
    "\n",
    "import onnx\n",
    "onnx_model = onnx.load(notebook_dir + \"/LFCW1A1.onnx\")\n",
    "from qonnx.core.modelwrapper import ModelWrapper\n",
    "onnx_model = ModelWrapper(onnx_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "showInNetron(notebook_dir + \"/LFCW1A1.onnx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qonnx.transformation.base import Transformation\n",
    "\n",
    "class ConvertSubToAdd(Transformation):\n",
    "    def apply(self, model):\n",
    "        graph = model.graph\n",
    "        for n in graph.node:\n",
    "            if n.op_type == \"Sub\":\n",
    "                A = model.get_initializer(n.input[1])\n",
    "                if A is not None:\n",
    "                    n.op_type = \"Add\"\n",
    "                    model.set_initializer(n.input[1], -A)\n",
    "        return (model, False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First the transformation class must be imported. Then a class can be created for the new transformation, which is derived from the base class. In this case the transformation has only the `apply()` function. \n",
    "\n",
    "All nodes are checked by first extracting the graph from the model and then iterating over the node list. With the help of .op_type the operation type of the node can be checked, if the node is a subtraction node the condition `if n.op_type == \"Sub\"` is true. It may be that the subtraction input of the node has no value, this is checked with `model.get_initializer(n.input[1])`. \n",
    "    \n",
    "    \n",
    "**Important:** FINN always assumes a certain order of inputs, this is especially important if you want to create additional custom operation nodes.\n",
    "\n",
    "When the input is initialized, the operation type of the node is converted to `\"Add\"`, this can simply be done by using the equal sign. Then the sign of the initial value must be changed. For this the ModelWrapper function `.set_initializer` can be used.\n",
    "\n",
    "At the end the changed model is returned and `model_was_changed` is set to False, because the transformation has to be executed only once.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "onnx_model_transformed = onnx_model.transform(ConvertSubToAdd())\n",
    "onnx_model_transformed.save('/tmp/LFCW1A1_changed.onnx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "showInNetron('/tmp/LFCW1A1_changed.onnx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parallel Transformation\n",
    "---------------------------------\n",
    "Some of the transformations in FINN can be performed in parallel on individual nodes. The following `NodeLocalTransformation` is required for this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qonnx.transformation.base import NodeLocalTransformation\n",
    "\n",
    "showSrc(NodeLocalTransformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transformations that are to be executed in parallel must have the method `applyNodeLocal()` implemented. Please note that the transformation is only executed on a single node, the parallel transformations do not have access to the entire model or tensors. Parallelization has the advantage that especially time-consuming transformations such as compilation can be executed more effectively. \n",
    "\n",
    "To control the degree of parallelization the argument `num_workers` can be specified. When the Docker container is started, the env variable `NUM_DEFAULT_WORKERS` is set to 4 by default, this can be increased or decreased depending on the system. You can also set the number of workers manually to a specific value when calling a transformation that allows parallelization. If the value is set to 0, all available CPU cores are used.\n",
    "\n",
    "In the following we want to take a closer look at the implementation using the compile transformation that is used for cpp simulation as example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from finn.transformation.fpgadataflow.compile_cppsim import CompileCppSim\n",
    "\n",
    "showSrc(CompileCppSim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The class is derived from the NodeLocalTransformation class and performs the compilation at every node that is an hls node."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
